#!/usr/bin/env python3
import cv2
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.callback_groups import ReentrantCallbackGroup
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from rcl_interfaces.msg import ParameterDescriptor, ParameterType
from geometry_msgs.msg import Point
from ball_collector_interfaces.srv import DetectTargetZone

# 导入MIPI相机模块
from hobot_vio import libsrcampy as srcampy

class TargetDetectionNode(Node):
    """
    目标区域检测节点，使用OpenCV识别紫色边框和目标颜色区域，直接使用MIPI相机
    """
    def __init__(self):
        super().__init__('target_detection_node')
        
        # 创建回调组
        self.service_cb_group = ReentrantCallbackGroup()
        
        # 声明参数
        self.declare_parameters(
            namespace='',
            parameters=[
                ('purple_hue_low', 140,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV色调下限')),
                ('purple_hue_high', 170,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV色调上限')),
                ('purple_sat_low', 50,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV饱和度下限')),
                ('purple_sat_high', 255,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV饱和度上限')),
                ('purple_val_low', 50,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV亮度下限')),
                ('purple_val_high', 255,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='紫色HSV亮度上限')),
                ('debug_view', True,
                 ParameterDescriptor(type=ParameterType.PARAMETER_BOOL,
                                     description='发布调试可视化结果')),
                ('camera_width', 640,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='相机宽度')),
                ('camera_height', 480,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='相机高度')),
                ('sensor_width', 1920,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='传感器宽度')),
                ('sensor_height', 1080,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='传感器高度'))
            ]
        )
        
        # 获取参数
        self.purple_hue_low = self.get_parameter('purple_hue_low').value
        self.purple_hue_high = self.get_parameter('purple_hue_high').value
        self.purple_sat_low = self.get_parameter('purple_sat_low').value
        self.purple_sat_high = self.get_parameter('purple_sat_high').value
        self.purple_val_low = self.get_parameter('purple_val_low').value
        self.purple_val_high = self.get_parameter('purple_val_high').value
        self.debug_view = self.get_parameter('debug_view').value
        self.camera_width = self.get_parameter('camera_width').value
        self.camera_height = self.get_parameter('camera_height').value
        self.sensor_width = self.get_parameter('sensor_width').value
        self.sensor_height = self.get_parameter('sensor_height').value
        
        # 颜色范围字典
        self.color_ranges = {
            'red': [
                (np.array([0, 100, 100]), np.array([10, 255, 255])),
                (np.array([160, 100, 100]), np.array([180, 255, 255]))  # 红色在HSV空间中分为两部分
            ],
            'green': [(np.array([40, 100, 100]), np.array([80, 255, 255]))],
            'blue': [(np.array([100, 100, 100]), np.array([140, 255, 255]))],
            'yellow': [(np.array([20, 100, 100]), np.array([30, 255, 255]))],
            'purple': [(np.array([140, 100, 100]), np.array([160, 255, 255]))],
            'orange': [(np.array([10, 100, 100]), np.array([20, 255, 255]))],
        }
        
        # 初始化CV桥接器
        self.bridge = CvBridge()
        
        # 存储最新图像
        self.current_frame = None
        self.frame_id = 'camera_frame'
        
        # 初始化MIPI相机
        self.init_camera()
        
        # 创建服务
        self.detect_service = self.create_service(
            DetectTargetZone,
            'detect_target_zone',
            self.detect_target_zone_callback,
            callback_group=self.service_cb_group
        )
        
        # 创建调试图像发布者
        if self.debug_view:
            self.debug_image_pub = self.create_publisher(
                Image,
                'target_detection_image',
                10
            )
            
            # 添加原始图像发布者，用于调试
            self.raw_image_pub = self.create_publisher(
                Image,
                'raw_image',
                10
            )
            
            self.get_logger().info('调试图像发布已启用')
            
        self.get_logger().info('目标区域检测节点已初始化')
    
    def init_camera(self):
        """初始化MIPI相机"""
        try:
            # 获取适合显示的分辨率
            self.disp_w, self.disp_h = self.get_display_res()
            
            # 创建Camera对象
            self.cam = srcampy.Camera()
            
            # 打开相机，使用与mipi_camera.py相同的参数
            self.cam.open_cam(0, -1, -1, 
                             [self.camera_width, self.disp_w], 
                             [self.camera_height, self.disp_h],
                             self.sensor_height, self.sensor_width)
            
            self.get_logger().info(f"MIPI相机初始化成功，分辨率: {self.camera_width}x{self.camera_height}")
            
            # 创建相机更新定时器
            self.camera_timer = self.create_timer(
                0.03,  # 约30FPS
                self.update_camera_frame
            )
            
        except Exception as e:
            self.get_logger().error(f"MIPI相机初始化失败: {e}")
            raise
    
    def get_display_res(self):
        """获取适合显示的分辨率，从mipi_camera.py移植"""
        disp_w_small = 640
        disp_h_small = 480
        
        try:
            disp = srcampy.Display()
            resolution_list = disp.get_display_res()
            
            if (self.sensor_width, self.sensor_height) in resolution_list:
                self.get_logger().info(f"分辨率 {self.sensor_width}x{self.sensor_height} 支持显示")
                return int(self.sensor_width), int(self.sensor_height)
            else:
                self.get_logger().info(f"分辨率 {self.sensor_width}x{self.sensor_height} 不在支持列表中，寻找合适的分辨率")
                for res in resolution_list:
                    # 排除0分辨率
                    if res[0] == 0 and res[1] == 0:
                        break
                    else:
                        disp_w_small = res[0]
                        disp_h_small = res[1]

                    # 找到不大于传感器分辨率的最大分辨率
                    if res[0] <= self.sensor_width and res[1] <= self.sensor_height:
                        self.get_logger().info(f"选择分辨率 {res[0]}x{res[1]}")
                        return int(res[0]), int(res[1])
            
            disp.close()
        except Exception as e:
            self.get_logger().error(f"获取显示分辨率失败: {e}")
        
        return disp_w_small, disp_h_small
    
    def update_camera_frame(self):
        """从MIPI相机获取图像帧"""
        try:
            # 获取相机图像
            img_data = self.cam.get_img(2, self.camera_width, self.camera_height)
            
            # 转换为numpy数组
            img_np = np.frombuffer(img_data, dtype=np.uint8)
            
            # 重塑为图像形状 (NV12格式)
            # NV12格式: 先是Y平面，后是交错的UV平面
            y_size = self.camera_width * self.camera_height
            y_plane = img_np[:y_size].reshape(self.camera_height, self.camera_width)
            
            # 创建与Y平面相同大小的U和V平面
            u_plane = np.zeros((self.camera_height, self.camera_width), dtype=np.uint8)
            v_plane = np.zeros((self.camera_height, self.camera_width), dtype=np.uint8)
            
            # 从交错的UV平面中提取U和V
            if len(img_np) > y_size:  # 确保有足够的数据包含UV平面
                uv_plane = img_np[y_size:].reshape(self.camera_height//2, self.camera_width//2, 2)
                
                # 放大UV平面到与Y平面相同的大小（简单的重复）
                u_plane[::2, ::2] = uv_plane[:, :, 0]
                u_plane[1::2, ::2] = uv_plane[:, :, 0]
                u_plane[::2, 1::2] = uv_plane[:, :, 0]
                u_plane[1::2, 1::2] = uv_plane[:, :, 0]
                
                v_plane[::2, ::2] = uv_plane[:, :, 1]
                v_plane[1::2, ::2] = uv_plane[:, :, 1]
                v_plane[::2, 1::2] = uv_plane[:, :, 1]
                v_plane[1::2, 1::2] = uv_plane[:, :, 1]
            
            # 合并YUV通道并转换为BGR格式
            yuv = cv2.merge([y_plane, u_plane, v_plane])
            bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
            
            # 更新当前帧
            self.current_frame = bgr
            
            # 发布原始图像用于调试
            if self.debug_view:
                try:
                    img_msg = self.bridge.cv2_to_imgmsg(bgr, "bgr8")
                    img_msg.header.stamp = self.get_clock().now().to_msg()
                    img_msg.header.frame_id = self.frame_id
                    self.raw_image_pub.publish(img_msg)
                except Exception as e:
                    self.get_logger().error(f"发布原始图像失败: {e}")
            
        except Exception as e:
            self.get_logger().error(f"获取相机图像失败: {e}")
    
    def detect_target_zone_callback(self, request, response):
        """检测目标区域服务回调"""
        if self.current_frame is None:
            self.get_logger().warn("没有可用的图像")
            response.success = False
            response.message = "没有可用的图像"
            return response
        
        try:
            # 检测目标区域
            found, position, border_point, size, debug_image = self.detect_target_zone(request.target_color)
            
            # 设置响应
            response.success = found
            response.message = "找到目标区域" if found else "未找到目标区域"
            
            if found:
                # 使用边界点而不是中心点，便于小车停在边缘
                response.center.x = float(border_point[0])
                response.center.y = float(border_point[1])
                response.size = float(size)
                self.get_logger().info(f"检测到目标区域: 位置=({border_point[0]}, {border_point[1]}), 尺寸={size}")
            else:
                self.get_logger().info(f"未检测到目标颜色 {request.target_color} 的目标区域")
            
            # 发布调试图像
            if self.debug_view and debug_image is not None:
                try:
                    img_msg = self.bridge.cv2_to_imgmsg(debug_image, "bgr8")
                    img_msg.header.stamp = self.get_clock().now().to_msg()
                    img_msg.header.frame_id = self.frame_id
                    self.debug_image_pub.publish(img_msg)
                except Exception as e:
                    self.get_logger().error(f"发布调试图像失败: {e}")
            
            return response
        
        except Exception as e:
            self.get_logger().error(f"检测目标区域失败: {e}")
            response.success = False
            response.message = f"检测失败: {str(e)}"
            return response
    
    def detect_target_zone(self, target_color):
        """
        检测目标区域(紫色边框+目标颜色内部)
        
        参数:
            target_color: 目标颜色
            
        返回:
            found: 是否找到目标区域
            position: 目标区域中心位置 (x, y)
            border_point: 紫色边框前方位置点(用于导航)
            size: 目标区域尺寸
            debug_image: 标记了检测结果的调试图像
        """
        # 复制图像以避免修改原始图像
        frame = self.current_frame.copy()
        debug_image = frame.copy() if self.debug_view else None
        
        # 转换到HSV颜色空间
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        
        # 紫色边框检测
        purple_lower = np.array([self.purple_hue_low, self.purple_sat_low, self.purple_val_low])
        purple_upper = np.array([self.purple_hue_high, self.purple_sat_high, self.purple_val_high])
        purple_mask = cv2.inRange(hsv, purple_lower, purple_upper)
        
        # 目标颜色检测 - 仅用于确认这是我们的目标区域
        target_mask = np.zeros_like(purple_mask)
        
        if target_color in self.color_ranges:
            for lower, upper in self.color_ranges[target_color]:
                color_part = cv2.inRange(hsv, lower, upper)
                target_mask = cv2.bitwise_or(target_mask, color_part)
        else:
            self.get_logger().warn(f"不支持的颜色: {target_color}")
            return False, (0, 0), (0, 0), 0, debug_image
        
        # 形态学操作以改善区域
        kernel = np.ones((5, 5), np.uint8)
        purple_mask = cv2.morphologyEx(purple_mask, cv2.MORPH_CLOSE, kernel)
        target_mask = cv2.morphologyEx(target_mask, cv2.MORPH_CLOSE, kernel)
        
        # 查找紫色边框轮廓
        purple_contours, _ = cv2.findContours(purple_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 查找目标颜色轮廓
        target_contours, _ = cv2.findContours(target_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 在调试图像中添加掩码
        if self.debug_view and debug_image is not None:
            # 显示紫色掩码
            purple_mask_small = cv2.resize(purple_mask, (self.camera_width//4, self.camera_height//4))
            purple_mask_color = cv2.cvtColor(purple_mask_small, cv2.COLOR_GRAY2BGR)
            debug_image[0:self.camera_height//4, 0:self.camera_width//4] = purple_mask_color
            
            # 显示目标颜色掩码
            target_mask_small = cv2.resize(target_mask, (self.camera_width//4, self.camera_height//4))
            target_mask_color = cv2.cvtColor(target_mask_small, cv2.COLOR_GRAY2BGR)
            debug_image[0:self.camera_height//4, self.camera_width//4:self.camera_width//2] = target_mask_color
            
            # 添加标签
            cv2.putText(debug_image, "紫色掩码", (10, 20), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.putText(debug_image, f"{target_color}掩码", (self.camera_width//4 + 10, 20), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        
        # 找到最大的紫色轮廓
        max_purple_cnt = None
        max_purple_area = 0
        
        for cnt in purple_contours:
            area = cv2.contourArea(cnt)
            if area > max_purple_area and area > 500:  # 设置最小面积阈值
                max_purple_area = area
                max_purple_cnt = cnt
        
        # 如果找到足够大的紫色轮廓
        if max_purple_cnt is not None:
            # 计算紫色边框的中心
            M = cv2.moments(max_purple_cnt)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                
                # 检查是否在紫色边框内有目标颜色
                has_target_color = False
                
                for target_cnt in target_contours:
                    # 检查目标颜色轮廓是否在紫色边框内
                    target_M = cv2.moments(target_cnt)
                    if target_M["m00"] != 0:
                        target_cx = int(target_M["m10"] / target_M["m00"])
                        target_cy = int(target_M["m01"] / target_M["m00"])
                        
                        # 简单检查: 判断目标颜色中心是否在紫色轮廓中
                        if cv2.pointPolygonTest(max_purple_cnt, (target_cx, target_cy), False) >= 0:
                            has_target_color = True
                            break
                
                # 如果找到目标区域(紫色边框+目标颜色)
                if has_target_color:
                    # 计算尺寸 (使用边界矩形的对角线长度)
                    x, y, w, h = cv2.boundingRect(max_purple_cnt)
                    size = np.sqrt(w**2 + h**2)
                    
                    # 计算边界点 - 这是小车应该停止的位置
                    # 选择靠近图像底部的紫色边框点
                    bottom_points = []
                    for point in max_purple_cnt:
                        px, py = point[0]
                        # 如果点在边框下半部分
                        if py > cy:
                            bottom_points.append((px, py))
                    
                    # 如果有下半部分的点，选择最接近相机中心x轴的点
                    border_point = (cx, cy)  # 默认使用中心点
                    if bottom_points:
                        # 根据与图像中心x轴的距离排序
                        center_x = self.camera_width / 2
                        bottom_points.sort(key=lambda p: abs(p[0] - center_x))
                        border_point = bottom_points[0]  # 最靠近中心x轴的点
                    
                    # 标记调试图像
                    if self.debug_view and debug_image is not None:
                        cv2.drawContours(debug_image, [max_purple_cnt], -1, (255, 0, 255), 2)
                        cv2.circle(debug_image, (cx, cy), 10, (0, 255, 0), -1)
                        cv2.circle(debug_image, border_point, 10, (0, 0, 255), -1)
                        cv2.putText(debug_image, "紫色边框", (cx - 60, cy - 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        cv2.putText(debug_image, "边界点", (border_point[0] - 60, border_point[1] + 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                        cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
                    
                    return True, (cx, cy), border_point, size, debug_image
        
        # 如果在调试模式下，标记所有轮廓
        if self.debug_view and debug_image is not None:
            # 显示紫色掩码轮廓
            for cnt in purple_contours:
                cv2.drawContours(debug_image, [cnt], -1, (255, 0, 255), 1)
            
            # 显示目标颜色掩码轮廓
            for cnt in target_contours:
                cv2.drawContours(debug_image, [cnt], -1, (0, 255, 0), 1)
            
            cv2.putText(debug_image, "未找到目标区域", (20, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        return False, (0, 0), (0, 0), 0, debug_image
    
    def destroy_node(self):
        """清理资源"""
        # 关闭相机
        if hasattr(self, 'cam'):
            try:
                self.cam.close_cam()
                self.get_logger().info("已关闭MIPI相机")
            except Exception as e:
                self.get_logger().error(f"关闭相机失败: {e}")
        
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)
    
    try:
        # 创建节点
        node = TargetDetectionNode()
        
        try:
            rclpy.spin(node)
        except KeyboardInterrupt:
            pass
        finally:
            node.destroy_node()
    except Exception as e:
        print(f"初始化节点失败: {e}")
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    main()